#include "Simulator.hpp"
#include "ParameterReader.hpp"
#include "VTKPathWriter.hpp"
#include <iostream>
#include <string.h>
#include <vector>
#include <stdlib.h>
#include <math.h>
#include <time.h>

Simulator::Simulator(const std::string& Paramfile, const std::string& ParticleFile)
{
    ParameterReader reader;

    if(!(reader.readParticleFile(masses,positions,velocities,ParticleFile,N_particles)
         && reader.readParameters(Paramfile)))
    {
        std::cerr << "ERROR: Exiting program"<<std::endl;
        exit(1);
    }

    reader.GetParameter("name",basename);
    reader.GetParameter("vis_space",vis_space);
    reader.GetParameter("t_start",t_start);
    reader.GetParameter("t_end",t_end);
    reader.GetParameter("delta_t",delta_t);
    reader.GetParameter("x_min",xmin);
    reader.GetParameter("y_min",ymin);
    reader.GetParameter("z_min",zmin);
    reader.GetParameter("x_max",xmax);
    reader.GetParameter("y_max",ymax);
    reader.GetParameter("z_max",zmax);
    reader.GetParameter("r_cut",r_cut);
    reader.GetParameter("epsilon",epsilon);
    reader.GetParameter("sigma",sigma);

    lengthx=xmax-xmin;
    lengthy=ymax-ymin;
    lengthz=zmax-zmin;
    nCx=floor(lengthx/r_cut);
    nCy=floor(lengthy/r_cut);
    nCz=floor(lengthz/r_cut);
    lenC_x=lengthx/nCx;
    lenC_y=lengthy/nCy;
    lenC_z=lengthz/nCz;
    N_cells=nCx*nCy*nCz;

    forces = new real4[N_particles];
    oldForces = new real4[N_particles];
    inverse2Masses = new real[N_particles];
    prodt2 = new real[N_particles];
    prodt = new real[N_particles];
    particle_ids = new int[N_particles];
    cells_ids=new int[N_cells];
    real dt2 = delta_t * delta_t;
    for (uint i = 0; i < N_particles; ++i)
    {
        particle_ids[i] = i;
        inverse2Masses[i]=1.0/(2.0*masses[i]);
        prodt2[i]=dt2*inverse2Masses[i];
        prodt[i]=delta_t*inverse2Masses[i];
        forces[i](0) = 0;
        forces[i](1) = 0;
        forces[i](2) = 0;
        oldForces[i](0) = 0;
        oldForces[i](1) = 0;
        oldForces[i](2) = 0;
    }

    setCellsToMinus();
    //Testing
    /* std::cout<<"The number of cells is: "<<N_cells<<"\n";
    std::cout<<nCx<<"  "<<nCy<<"  "<<nCy<<"\n";
    std::cout<<"Particle  Mass  PositionX  PositionY PositionZ PositionX  VelocityX  VelocityY  VelocityZ   ForceX  ForceY  ForceZ\n";
    for(uint i=0;i<N_particles;++i)
    {
        std::cout<<i<<" "<<masses[i]<<"  "<<inverse2Masses[i]<<"  "<<positions[i](0)<<" "<<" "<<positions[i](1)<<" "<<positions[i](2)
                <<" "<<velocities[i](0)<<" "<<velocities[i](1)<<" "<<velocities[i](2)<<" "<<forces[i](0)<<" "<<forces[i](1)<<" "<<forces[i](2)<<"\n";
    }
   */
}



void Simulator::updatePositions()
{
    for(uint i=0;i<N_particles;++i)
    {
        positions[i]=positions[i]+delta_t*velocities[i]+prodt2[i]*forces[i];
        oldForces[i]=forces[i];
    }

}

void Simulator::updateVelocities()
{
    for(uint i=0;i<N_particles;++i)
    {
        velocities[i]+=(forces[i]+oldForces[i])*prodt[i];
    }
}

void Simulator::assembleCells()
{
    uint cx,cy,cz, cellID;
    real inv_x = 1.0/lenC_x;
    real inv_y = 1.0/lenC_y;
    real inv_z = 1.0/lenC_z;

    for(uint i=0;i<N_particles;++i)
    {
        cx=positions[i](0)*inv_x;
        cy=positions[i](1)*inv_y;
        cz=positions[i](2)*inv_z;

        cellID=cx*nCy*nCz+cy*nCz+cz;

        particle_ids[i]=cells_ids[cellID];
        cells_ids[cellID]=i;
    }
}

void Simulator::calcForces(const uint& fXboundaries,const uint& fYboundaries,const uint& fZboundaries)
{
    uint cx,cy,cz,cellID;
    uint neighbor_cx, neighbor_cy, neighbor_cz;
    real4 sumForce;
    real sma6 = sigma * sigma * sigma;
    sma6*=sma6;
    real inv_x = 1.0/lenC_x;
    real inv_y = 1.0/lenC_y;
    real inv_z = 1.0/lenC_z;
    for(uint i=0;i<N_particles;++i)
    {
        cx=positions[i](0)*inv_x;
        cy=positions[i](1)*inv_y;
        cz=positions[i](2)*inv_z;

        sumForce(0)=0.0;
        sumForce(1)=0.0;
        sumForce(2)=0.0;
        sumForce(3)=0.0;
        for(int l=-1;l<=1;++l)
        {
            for(int m=-1;m<=1;++m)
            {
                for(int n=-1;n<=1;++n)
                {
                    neighbor_cx=cx+l;
                    neighbor_cy=cy+m;
                    neighbor_cz=cz+n;

                    int flagX=0;
                    int flagY=0;
                    int flagZ=0;

                    if(fXboundaries==PERIODIC)
                    {
                        if(neighbor_cx<0)
                        {
                            neighbor_cx=nCx-1;
                            flagX=1;
                        }
                        else if(neighbor_cx>nCx-1)
                        {
                            neighbor_cx=0;
                            flagX=1;
                        }
                    }
                    else
                    {
                        if(neighbor_cx<0 || neighbor_cx>nCx-1)
                            continue;
                    }

                    if(fYboundaries==PERIODIC)
                    {
                        if(neighbor_cy<0)
                            neighbor_cy=nCy-1;
                        else if(neighbor_cy>nCy-1)
                            neighbor_cy=0;
                    }
                    else
                    {
                        if(neighbor_cy<0 || neighbor_cy>nCy-1)
                            continue;
                    }

                    if(fZboundaries==PERIODIC)
                    {
                        if(neighbor_cz<0)
                        {
                            neighbor_cz=nCz-1;
                            flagY=1;
                        }
                        else if(neighbor_cz>nCz-1)
                        {
                            neighbor_cz=0;
                            flagZ=1;
                        }
                    }
                    else
                    {
                        if(neighbor_cz<0 || neighbor_cz>nCz-1)
                            continue;
                    }
                    cellID=neighbor_cx*nCy*nCz+neighbor_cy*nCz+neighbor_cz;
                    int j=cells_ids[cellID];
                    while(j!=-1)
                    {
                        real4 dij = positions[j]-positions[i];
                        if(flagX==1)
                            dij(0)=dij(0)-round(dij(0)*nCx*inv_x);
                        if(flagY==1)
                            dij(1)=dij(1)-round(dij(1)*nCy*inv_y);
                        if(flagZ==1)
                            dij(2)=dij(2)-round(dij(2)*nCz*inv_z);
                        dij(3)=0.0;
                        real n2sqr = dij(0)*dij(0) + dij(1)*dij(1) + dij(2)*dij(2);
                        if(n2sqr<=r_cut*r_cut)
                        {
                            if (static_cast<int>(i) != j)
                            {
                                real num = sma6*(n2sqr*n2sqr*n2sqr-2.0*sma6);
                                real den = n2sqr*n2sqr*n2sqr*n2sqr*n2sqr*n2sqr*n2sqr;
                                real expression = num/den;
                                dij = dij*expression;
                                sumForce+=dij;
                            }
                        }
                        j=particle_ids[j];
                    }
                    forces[i]=24*epsilon*sumForce;
                }
            }
        }

    }
}

void Simulator::setBoundaries(const uint& fXboundaries,const uint& fYboundaries,const uint& fZboundaries)
{
    for(uint i=0;i<N_particles;++i)
    {
        if(fXboundaries==PERIODIC)
        {
            if(positions[i](0)<xmin)
                positions[i](0)+=lengthx;
            else if(positions[i](0)>xmax)
                positions[i](0)-=lengthx;
        }
        else
        {
            //TODO
        }
        if(fYboundaries==PERIODIC)
        {
            if(positions[i](1)<ymin)
                positions[i](1)+=lengthy;
            else if(positions[i](1)>ymax)
                positions[i](1)-=lengthy;
        }
        else
        {
            //TODO
        }
        if(fZboundaries==PERIODIC)
        {
            if(positions[i](2)<zmin)
                positions[i](2)+=lengthz;
            else if(positions[i](2)>zmax)
                positions[i](2)-=lengthz;
        }
        else
        {
            //TODO
        }
    }
}

void Simulator::sim(const uint& fXboundaries,const uint& fYboundaries,const uint& fZboundaries)
{
    std::string filename;
    int steps = (t_end-t_start)/delta_t;
    int n=1;

    //write vtk file before simulation
    filename = "VTK/" + basename;
    filename+=std::to_string(n-1);
    VTKPathWriter vtkWriter(filename,N_particles);
    vtkWriter.writePoints(positions);
    vtkWriter.writeScalars(masses,"mass");
    vtkWriter.writeVector(forces,"force");
    vtkWriter.writeVector(velocities,"velocity");

    //Initial forces
    assembleCells();
    calcForces(fXboundaries,fYboundaries,fZboundaries);

    for(int t = 0;t<=steps;++t)
    {
        if (t!=0 && (t%vis_space) == 0)
        {
            filename = "VTK/" + basename;
            filename+=std::to_string(n);
            VTKPathWriter vtkWriter(filename,N_particles);
            vtkWriter.writePoints(positions);
            vtkWriter.writeScalars(masses,"mass");
            vtkWriter.writeVector(forces,"force");
            vtkWriter.writeVector(velocities,"velocity");

            ++n;
        }
        updatePositions();
        setBoundaries(fXboundaries,fYboundaries,fZboundaries);
        setCellsToMinus();
        assembleCells();
        calcForces(fXboundaries,fYboundaries,fZboundaries);
        updateVelocities();

    }
}
